from diffusers import StableDiffusion3Pipeline
import torch

def text2image(model_name, prompt, num_inference_steps, guidance_scale, output_path):
    pipe = StableDiffusion3Pipeline.from_pretrained(model_name,
                                                    torch_dtype=torch.float16)
    pipe = pipe.to("cuda")

    image = pipe(
        prompt,
        negative_prompt="",
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
    ).images[0]
    image.save(output_path)


def main():
    prompt = "A cinematic shot of a baby racoon wearing an intricate italian priest robe."
    model_name = "/vepfs_train2/meichaoyang/model/stabilityai/stable-diffusion-3-medium-diffusers"
    text2image(model_name, prompt, 28, 7, "baby_racoon.png")

if __name__ == '__main__':
    main()
